#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能体任务编排 - Streamlit Web界面
提供用户友好的任务执行界面，支持实时消息显示
"""
import streamlit as st
import asyncio
import threading
import time
import queue
import json
from datetime import datetime
from typing import Dict, Any, Optional
import sys
import os

# 添加项目根目录到路径
sys.path.append(os.path.dirname(__file__))

from src.task3.task3_main import TaskOrchestrator3


class MessageQueue:
    """消息队列管理器，用于后台服务与前端的通信"""
    
    def __init__(self):
        self.message_queue = queue.Queue()
        self.is_running = False
    
    def add_message(self, message: str, message_type: str = "info", timestamp: Optional[str] = None):
        """添加消息到队列"""
        if timestamp is None:
            timestamp = datetime.now().strftime("%H:%M:%S")
        
        msg_data = {
            "content": message,
            "type": message_type,  # info, success, error, warning
            "timestamp": timestamp
        }
        self.message_queue.put(msg_data)
    
    def get_messages(self):
        """获取所有待处理的消息"""
        messages = []
        while not self.message_queue.empty():
            try:
                messages.append(self.message_queue.get_nowait())
            except queue.Empty:
                break
        return messages
    
    def clear(self):
        """清空消息队列"""
        while not self.message_queue.empty():
            try:
                self.message_queue.get_nowait()
            except queue.Empty:
                break


class StreamlitTaskOrchestrator(TaskOrchestrator3):
    """继承TaskOrchestrator3，添加Streamlit消息输出功能"""
    
    def __init__(self, message_queue: MessageQueue):
        super().__init__()
        self.message_queue = message_queue
    
    def log_message(self, message: str, msg_type: str = "info"):
        """发送消息到Streamlit界面"""
        self.message_queue.add_message(message, msg_type)
    
    async def execute_task(self, topic: str) -> None:
        """重写执行任务方法，添加消息输出"""
        self.log_message(f"🚀 开始执行报告生成任务：{topic}", "info")
        
        try:
            # 调用父类方法
            await super().execute_task(topic)
            self.log_message("✅ 所有任务执行完成！", "success")
            self.log_message(f"📄 报告框架文件：report_framework.md", "success")
            
        except Exception as e:
            self.log_message(f"❌ 任务执行失败：{str(e)}", "error")
            raise


def init_session_state():
    """初始化Session State"""
    if 'message_queue' not in st.session_state:
        st.session_state.message_queue = MessageQueue()
    
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    
    if 'task_running' not in st.session_state:
        st.session_state.task_running = False
    
    if 'orchestrator' not in st.session_state:
        st.session_state.orchestrator = None


def display_chat_message(message: Dict[str, Any]):
    """显示单条聊天消息"""
    timestamp = message.get("timestamp", "")
    content = message.get("content", "")
    msg_type = message.get("type", "info")
    
    # 根据消息类型选择图标和颜色
    type_config = {
        "info": {"icon": "💬", "color": "#1f77b4"},
        "success": {"icon": "✅", "color": "#2ca02c"},
        "error": {"icon": "❌", "color": "#d62728"},
        "warning": {"icon": "⚠️", "color": "#ff7f0e"},
        "system": {"icon": "🤖", "color": "#9467bd"}
    }
    
    config = type_config.get(msg_type, type_config["info"])
    
    with st.container():
        col1, col2 = st.columns([0.1, 0.9])
        with col1:
            st.markdown(f"<div style='font-size: 20px;'>{config['icon']}</div>", 
                       unsafe_allow_html=True)
        with col2:
            st.markdown(f"""
            <div style='
                background-color: rgba({int(config['color'][1:3], 16)}, {int(config['color'][3:5], 16)}, {int(config['color'][5:7], 16)}, 0.1);
                padding: 10px;
                border-radius: 10px;
                border-left: 4px solid {config['color']};
                margin-bottom: 10px;
            '>
                <div style='font-size: 12px; color: #666; margin-bottom: 5px;'>{timestamp}</div>
                <div style='color: #333;'>{content}</div>
            </div>
            """, unsafe_allow_html=True)


def update_chat_display():
    """更新聊天显示区域"""
    # 获取新消息
    new_messages = st.session_state.message_queue.get_messages()
    
    # 添加到历史记录
    for msg in new_messages:
        st.session_state.chat_history.append(msg)
    
    # 限制历史记录长度
    if len(st.session_state.chat_history) > 100:
        st.session_state.chat_history = st.session_state.chat_history[-100:]


async def run_task_async(topic: str, message_queue: MessageQueue):
    """异步运行任务"""
    try:
        orchestrator = StreamlitTaskOrchestrator(message_queue)
        await orchestrator.execute_task(topic)
    except Exception as e:
        message_queue.add_message(f"任务执行异常：{str(e)}", "error")


def run_task_in_thread(topic: str, message_queue: MessageQueue):
    """在线程中运行异步任务"""
    try:
        # 创建新的事件循环
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        # 运行任务
        loop.run_until_complete(run_task_async(topic, message_queue))
        
    except Exception as e:
        message_queue.add_message(f"线程执行异常：{str(e)}", "error")
    finally:
        message_queue.add_message("🏁 任务线程结束", "system")
        # 重置运行状态
        st.session_state.task_running = False


def main():
    """主函数"""
    # 页面配置
    st.set_page_config(
        page_title="智能体任务编排系统",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # 初始化
    init_session_state()
    
    # 页面标题
    st.title("🤖 智能体任务编排系统")
    st.markdown("---")
    
    # 侧边栏 - 任务配置
    with st.sidebar:
        st.header("📋 任务配置")
        
        # 主题输入
        topic = st.text_input(
            "请输入报告主题：",
            placeholder="例如：2025 Q3 AI大模型安全政策法规",
            help="输入您想要生成报告的主题"
        )
        
        # 执行按钮
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🚀 开始执行", disabled=st.session_state.task_running):
                if topic.strip():
                    st.session_state.task_running = True
                    st.session_state.message_queue.add_message(f"📝 收到任务请求：{topic}", "system")
                    
                    # 在后台线程中执行任务
                    thread = threading.Thread(
                        target=run_task_in_thread,
                        args=(topic, st.session_state.message_queue),
                        daemon=True
                    )
                    thread.start()
                    
                    st.success("任务已启动！")
                else:
                    st.error("请输入有效的主题！")
        
        with col2:
            if st.button("🗑️ 清空对话"):
                st.session_state.chat_history.clear()
                st.session_state.message_queue.clear()
                st.success("对话已清空！")
        
        # 任务状态
        st.markdown("---")
        st.subheader("📊 任务状态")
        if st.session_state.task_running:
            st.success("🟢 任务运行中...")
        else:
            st.info("🔵 等待任务...")
        
        # 统计信息
        st.markdown("---")
        st.subheader("📈 统计信息")
        st.metric("消息总数", len(st.session_state.chat_history))
        
        # 系统信息
        st.markdown("---")
        st.subheader("ℹ️ 系统信息")
        st.info("""
        **功能说明：**
        - 输入主题后点击"开始执行"
        - 系统将自动执行智能体任务编排
        - 实时显示任务执行进度
        - 支持后台服务消息推送
        """)
    
    # 主区域 - 对话显示
    st.header("💬 任务执行对话")
    
    # 创建对话容器
    chat_container = st.container()
    
    # 自动刷新机制
    if st.session_state.task_running:
        # 任务运行时每秒刷新
        time.sleep(1)
        st.rerun()
    else:
        # 任务未运行时每3秒检查一次
        time.sleep(3)
        st.rerun()
    
    # 更新聊天显示
    update_chat_display()
    
    # 显示聊天历史
    with chat_container:
        if st.session_state.chat_history:
            # 创建可滚动的对话区域
            with st.container():
                for message in st.session_state.chat_history:
                    display_chat_message(message)
        else:
            st.info("👋 欢迎使用智能体任务编排系统！请在左侧输入主题开始任务。")
    
    # 底部信息
    st.markdown("---")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.markdown("**🔧 开发框架：** AgentScope")
    with col2:
        st.markdown("**🌐 Web框架：** Streamlit")
    with col3:
        st.markdown("**🤖 AI模型：** 阿里云通义千问")


# 外部接口函数，供其他服务调用
def send_message_to_streamlit(message: str, msg_type: str = "info"):
    """
    外部接口：向Streamlit页面发送消息
    
    Args:
        message: 消息内容
        msg_type: 消息类型 (info, success, error, warning, system)
    
    Usage:
        from streamlit_web import send_message_to_streamlit
        send_message_to_streamlit("Hello from backend!", "info")
    """
    if 'message_queue' in st.session_state:
        st.session_state.message_queue.add_message(message, msg_type)


def get_task_status():
    """
    外部接口：获取当前任务状态
    
    Returns:
        dict: 包含任务状态信息的字典
    """
    if hasattr(st.session_state, 'task_running'):
        return {
            "running": st.session_state.task_running,
            "message_count": len(st.session_state.chat_history) if hasattr(st.session_state, 'chat_history') else 0
        }
    return {"running": False, "message_count": 0}


if __name__ == "__main__":
    main()